from torch import nn
import torch
from typing import Tuple


class Inception(nn.Module):
    # `c1`--`c4` 是每条路径的输出通道数
    def __init__(self, in_channels: int, c1: int, c2: Tuple[int, int], c3: Tuple[int, int], c4: int):
        super(Inception, self).__init__()
        self._act = nn.ReLU(inplace=True)
        # 线路1,单1 x 1卷积层
        self._p1_1 = nn.Conv2d(in_channels, c1, kernel_size=(1,))
        # 线路2,1 x 1卷积层后接3 x 3卷积层
        self._p2_1 = nn.Conv2d(in_channels, c2[0], kernel_size=(1,))
        self._p2_2 = nn.Conv2d(c2[0], c2[1], kernel_size=(3,), padding=(1,))
        # 线路3,1 x 1卷积层后接5 x 5卷积层
        self._p3_1 = nn.Conv2d(in_channels, c3[0], kernel_size=(1,))
        self._p3_2 = nn.Conv2d(c3[0], c3[1], kernel_size=(5,), padding=2)
        # 线路4,3 x 3最大汇聚层后接1 x 1卷积层
        self._p4_1 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
        self._p4_2 = nn.Conv2d(in_channels, c4, kernel_size=(1,))

    def forward(self, x):
        p1 = self._act(self._p1_1(x))
        p2 = self._act(self._p2_2(self._act(self._p2_1(x))))
        p3 = self._act(self._p3_2(self._act(self._p3_1(x))))
        p4 = self._act(self._p4_2(self._p4_1(x)))
        return torch.cat((p1, p2, p3, p4), dim=1)


class GoogLeNet(nn.Module):
    def __init__(self):
        super(GoogLeNet, self).__init__()
        self._stage_1 = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=(7,), stride=(2,), padding=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )
        self._stage_2 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=(1,)),
            nn.ReLU(),
            nn.Conv2d(64, 192, kernel_size=(3,), padding=1),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )
        self._stage_3 = nn.Sequential(
            Inception(192, 64, (96, 128), (16, 32), 32),
            Inception(256, 128, (128, 192), (32, 96), 64),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )
        self._stage_4 = nn.Sequential(
            Inception(480, 192, (96, 208), (16, 48), 64),
            Inception(512, 160, (112, 224), (24, 64), 64),
            Inception(512, 128, (128, 256), (24, 64), 64),
            Inception(512, 112, (144, 288), (32, 64), 64),
            Inception(528, 256, (160, 320), (32, 128), 128),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )
        self._stage_5 = nn.Sequential(
            Inception(832, 256, (160, 320), (32, 128), 128),
            Inception(832, 384, (192, 384), (48, 128), 128),
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten()
        )
        self._linear = nn.Linear(1024, 10)

    def forward(self, x):
        x = self._stage_1(x)
        x = self._stage_2(x)
        x = self._stage_3(x)
        x = self._stage_4(x)
        x = self._stage_5(x)
        return self._linear(x)
